import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F


class cnnEncoder(nn.Module):
    def __init__(self):
        super(cnnEncoder, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1) 
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))     
        x = F.max_pool2d(x, 2, 2)   
        x = F.relu(self.conv2(x))    
        x = F.max_pool2d(x, 2, 2)   
        x = x.view(-1, 4*4*50)      
        x = F.relu(self.fc1(x))   
        x = self.fc2(x)  

        return x


def init_weight(layer):
    if type(layer) == nn.Conv2d:
        nn.init.constant_(layer.weight, 1)
    elif type(layer) == nn.Linear:
        nn.init.constant_(layer.weight, 1)
        nn.init.constant_(layer.bias, 0)


if __name__ == "__main__":
    encoder = cnnEncoder()
    encoder.apply(init_weight)
    # print(encoder.fc.weight)
    # print(encoder.fc.bias)

    x = np.random.randn(1, 1, 28, 28).astype(np.float32)
    x = torch.FloatTensor(x)
    print(x)
    result = encoder(x)
    print(result)
    # print(torch.mean(result))


        